from abc import ABC

import gym
import numpy as np
from gym import spaces

from pb.ros.geometry_msgs import Pose_pb2
from pb.ros.geometry_msgs import PoseStamped_pb2
from pb.ros.nav_msgs import Odometry_pb2
from pb.ros.visualization_msgs import Marker_pb2
from pb.ros.geometry_msgs import Point_pb2
from pb.ros.std_msgs import Bool_pb2
from pb.ros.sensor_msgs import Image_pb2
import proxy

import os
import time
import math
import threading
import sys

import logging

from uav_gym.utils.convert_angle import quaternion_to_euler, euler_to_quaternion

node = proxy.Connector('tcp://172.22.56.12:15555')  # 在709的电脑上使用该IP
# node = proxy.Connector('tcp://12.12.12.204:15555') #在集群的电脑上使用该IP



class RotodyneEnv(gym.Env):
    metadata = {}

    # TODO
    def __init__(self, env_config):
        # NotNeedTODO: define action space. George has set action_bound in the learner.
        # action: delta_position which is relative to the current odom position [x,y,z,yaw]
        self.frame_skip = env_config['frame_skip']

        self.odom_position = None
        self.odom_vel = None
        self.depth_image = None
        self.kinematics_desired_position_quaternion = None
        self.goal = None  # [x,y,z,yaw]

        # Initialize ROS
        self.subcribe_odom_position = node.Subscribe('FK', 'M300_0', '/FK/M300_0/odom', Odometry_pb2.Odometry,
                                                self.callback_odom_position_velocity)  # 里程计---相对于起点位置
        self.subcribe_kinematics_desired_position_quaternion = node.Subscribe('FK', 'M300_0', '/FK/M300_0/desired_pos', Pose_pb2.Pose,
                                              self.callback_kinematics_desired_position_quaternion)  # 期望位置控制指令(浙大控制算法)---相对于起点位置
        self.subscribe_depth_image = node.Subscribe('FK', 'M300_0', '/FK/M300_0/depth/image_raw', Image_pb2.Image, self.callback_depth_image)



        self.publish_restart = node.Advertise('FK', 'System', '/FK/reset', Bool_pb2.Bool)

        self.publish_goal = node.Advertise('FK', 'M300_0', '/FK/M300_0/cmd_goal',
                                          PoseStamped_pb2.PoseStamped)  # 设置终点（相对位置）
        self.publish_dynamics_desired_position_quaternion  = node.Advertise('FK', 'M300_0', '/FK/M300_0/cmd_pos', Pose_pb2.Pose)  # 发给动力学模型的控制输入

        self.publish_stepforward = node.Advertise('FK', 'System', '/FK/resume', Bool_pb2.Bool)


        self.reset()


    def callback_odom_position_velocity(self, odom):
        x2 = odom.pose.pose.position.x
        y2 = odom.pose.pose.position.y
        z2 = odom.pose.pose.position.z
        x = odom.pose.pose.orientation.x
        y = odom.pose.pose.orientation.y
        z = odom.pose.pose.orientation.z
        w = odom.pose.pose.orientation.w
        self.odom_position = np.array([x2, y2, z2, x, y, z, w],dtype=np.float32)

        angular_x = odom.twist.twist.angular.x
        angular_y = odom.twist.twist.angular.y
        angular_z = odom.twist.twist.angular.z
        linear_x = odom.twist.twist.linear.x
        linear_y = odom.twist.twist.linear.y
        linear_z = odom.twist.twist.linear.z
        self.odom_vel = np.array([angular_x, angular_y, angular_z, linear_x, linear_y, linear_z],dtype=np.float32)

    def callback_kinematics_desired_position_quaternion(self, pos):
        x1 = pos.position.x
        y1 = pos.position.y
        z1 = pos.position.z
        x = pos.orientation.x
        y = pos.orientation.y
        z = pos.orientation.z
        w = pos.orientation.w
        self.kinematics_desired_position_quaternion = np.array([x1,y1,z1,x,y,z,w])
        # print("desired_pos:", x1, '\n')

    def callback_depth_image(self, image):
        imagedata = np.frombuffer(image.data, np.uint16)
        self.depth_image = imagedata.reshape(1, image.height, image.width) # channel is 1
        self.depth_image = self.depth_image.astype(np.float32)    # millimeter -> meter

    def set_goal(self,goal):
        cmd_goal = PoseStamped_pb2.PoseStamped()
        cmd_goal.pose.position.x = goal[0]
        cmd_goal.pose.position.y = goal[1]
        cmd_goal.pose.position.z = goal[2]
        cmd_goal.pose.orientation.x = goal[3]
        cmd_goal.pose.orientation.y = goal[4]
        cmd_goal.pose.orientation.z = goal[5]
        cmd_goal.pose.orientation.w = goal[6]
        cmd_goal.header.frame_id = "world"
        self.publish_goal.Publish(cmd_goal)

    def set_dynamics_desired_position(self,dynamics_desired_position_quaternion):
        dynamics_desired_position_quaternion_ros = Pose_pb2.Pose()
        dynamics_desired_position_quaternion_ros.position.x = dynamics_desired_position_quaternion[0]
        dynamics_desired_position_quaternion_ros.position.y = dynamics_desired_position_quaternion[1]
        dynamics_desired_position_quaternion_ros.position.z = dynamics_desired_position_quaternion[2]
        dynamics_desired_position_quaternion_ros.orientation.x = dynamics_desired_position_quaternion[3]
        dynamics_desired_position_quaternion_ros.orientation.y = dynamics_desired_position_quaternion[4]
        dynamics_desired_position_quaternion_ros.orientation.z = dynamics_desired_position_quaternion[5]
        dynamics_desired_position_quaternion_ros.orientation.w = dynamics_desired_position_quaternion[6]
        self.publish_dynamics_desired_position_quaternion.Publish(dynamics_desired_position_quaternion_ros)

    def step_forward(self):
        dbool = Bool_pb2.Bool()
        dbool.data = True
        self.publish_stepforward.Publish(dbool)



    def get_obs(self):
        # TODO: Depth Image Observation

        # position observation
        self.subcribe_odom_position.SpinOnce()
        self.subcribe_kinematics_desired_position_quaternion.SpinOnce()
        self.subscribe_depth_image.SpinOnce()
        return {"depth_image":self.depth_image, "odom_vel": self.odom_vel, "odom_position": self.odom_position,
                "kinematics_desired_position_quaternion": self.kinematics_desired_position_quaternion, "goal":self.goal}

    def reset(self, seed=None, options=None):

        # kong has fixed the reset bug
        # reset the simulation
        abool = Bool_pb2.Bool()
        abool.data = True
        self.publish_restart.Publish(abool)
        # logging.warning("Not reset the simulation due to the Kong bug")

        self.generate_goal_point()
        # reset the desired_position to 0
        self.set_goal(self.goal)

        # call step_forward() to prove there exists messages in ROS
        self.step_forward()

        # NoNeedTODO: reset the kinematics-algorithm, Kong says it is no need to reset kinematics-algorithm on the beginning of new episodes

    def generate_goal_point(self):
        # TODO
        self.goal = np.array([0,0,0,0,0,0,1],dtype=np.float32)


    # TODO
    def step(self, action):
        obs = self.get_obs()
        self.do_simulation(obs, action, self.frame_skip)

        # reward calculation
        # TODO: reward1: using kinematics guidance
        # kinematics_desired_position_euler = quaternion_to_euler(self.kinematics_desired_position_quaternion)
        reward1 = 0
        # TODO: reward2: distance to the goal
        reward2 = 0
        # TODO: reward3: crash
        reward3 = 0
        reward = reward1 + reward2 + reward3

        # TODO: termination condition
        terminated = False
        next_obs = self.get_obs()

        # TODO: 'Reach Goal', 'Crashed', or 'Time out'
        status = "Not Implemented"
        return reward, terminated, next_obs, status

    def do_simulation(self, obs, action, frame_skip=1):
        # kinematics_desired_position_euler = quaternion_to_euler(self.kinematics_desired_position_quaternion)
        # obs_euler = quaternion_to_euler([obs["odom_position"][3], obs["odom_position"][4], obs["odom_position"][5],
        #                     obs["odom_position"][6]])
        # dynamics_desired_position_euler = action + [obs["odom_position"][0],obs["odom_position"][1],obs["odom_position"][2],obs_euler[-1]]
        # quaternion =
        # dynamics_desired_position_quaternion = [dynamics_desired_position_euler[0],dynamics_desired_position_euler[2],dynamics_desired_position_euler[3]]
        # self.set_dynamics_desired_position(dynamics_desired_position_quaternion)

        _,_,obs_yaw = quaternion_to_euler([obs["odom_position"][3], obs["odom_position"][4], obs["odom_position"][5],
                            obs["odom_position"][6]])
        desired_yaw = obs_yaw + action[-1]
        desired_quaternion = euler_to_quaternion(0,0,desired_yaw)
        dynamics_desired_position_quaternion = np.array([obs["odom_position"][0]+action[0],obs["odom_position"][1]+action[1],obs["odom_position"][2]+action[2],
                                                desired_quaternion[0],desired_quaternion[1],desired_quaternion[2],desired_quaternion[3]])
        self.set_dynamics_desired_position(dynamics_desired_position_quaternion)

        # from Jiaqing code
        # Step forward environment and dynamics
        # TODO: Check whether step forward succeed
        for _ in range(frame_skip):
            self.step_forward()




